import datetime
import os
import sys
import time
import torch
import tqdm

from common_runner import put_data_to_device
from ai.run.base_train_runner import GxlBaseRunner


class RunnerGxl(GxlBaseRunner):
    def __init__(self, model, optim, loss_f, train_loader,
                 config, valid_loader=None, scheduler=None, multi=False,
                 local_rank=0, is_class=True,
                 device=torch.device('cpu'), model_save_dir='./output/model2/'):

        super(RunnerGxl, self).__init__(model, optim, loss_f, train_loader,
                                        config, valid_loader, scheduler, multi,
                                        local_rank, is_class, device)
        os.makedirs(model_save_dir, exist_ok=True)
        self.model_save_dir = model_save_dir

    def calculate_valid_loss(self, ):
        self.model.eval()
        all_accuracy = 0.0
        batch_size = 0
        with torch.no_grad():
            for X, l, y in self.valid_loader:
                X, y = put_data_to_device(X, y, self.device)
                out = self.model(X, )
                accuracy = self.accuracy(out, y)
                all_accuracy += accuracy
                batch_size += 1
        return all_accuracy / batch_size

    def calculate_valid_accuracy(self):
        self.model.eval()
        all_accuracy = 0.0
        batch_num = 0
        with torch.no_grad():
            for X,l, y in self.valid_loader:
                X, y = put_data_to_device(X, y, self.device)
                out = self.model(X, )
                accuracy = self.accuracy(out, y)
                all_accuracy += accuracy
                batch_num += 1
        return all_accuracy / batch_num

    def handle_batch(self, batch):
        comments,l, titles = [x.to(self.device) for x in batch]
        Y_hat = self.model(comments)
        l = self.loss_f(Y_hat, titles)
        l.sum().backward()  # 损失函数的标量进⾏“反向传播”
        # utils.utils_model.grad_clipping(self.model, 1)
        self.optim.step()
        if self.scheduler:
            self.scheduler.step()
        return l.sum().item()
